#ifdef __CUDACC__
extern "C"{
#endif

#include <cutil_inline.h>
#include <limits.h>

#include "cuda-me.h"

#define SEARCHWINDOW 8
#define LAMBDA 4
#define TILE 16

// declare texture reference for 2D float texture
texture<uint8_t, 2, cudaReadModeElementType> currentTex;
texture<uint8_t, 2, cudaReadModeElementType> lastTex;

__global__ void minSAD(unsigned bw, int* mvXp, int* mvYp, int* mvX, int* mvY, int* sad);

void runCuda(int width, int height,  int* mvXp, int* mvYp, int* mvX, int* mvY, int* sad);

int* mvX = NULL;
int* mvY = NULL;
int* mvXp = NULL;
int* mvYp = NULL;
int* sad = NULL;

cudaChannelFormatDesc channelDesc;

cudaArray* currentArray, *lastArray;

void init_cuda(int i_width, int i_lines)
{
    currentTex.addressMode[0] = cudaAddressModeClamp;
    lastTex.addressMode[0] = cudaAddressModeClamp;

	int bw = (i_width + 15) / 16;
	int bh = (i_lines + 15) / 16;

    channelDesc = cudaCreateChannelDesc<uint8_t>();

	cutilSafeCall(cudaMallocArray(&currentArray, 	&channelDesc, i_width,	i_lines));
	cutilSafeCall(cudaMallocArray(&lastArray, 		&channelDesc, i_width, 	i_lines));

    cudaMalloc((void**)&sad, bw * bh * sizeof(int));

    cudaMalloc((void**)&mvX, bw * bh * sizeof(int));
    cudaMalloc((void**)&mvY, bw * bh * sizeof(int));

    cudaMalloc((void**)&mvXp, bw * bh * sizeof(int));
    cudaMalloc((void**)&mvYp, bw * bh * sizeof(int));
}

void deinit_cuda()
{
    cutilSafeCall(cudaFree(mvX));
    cutilSafeCall(cudaFree(mvY));

	cutilSafeCall(cudaFree(mvXp));
    cutilSafeCall(cudaFree(mvYp));

    cutilSafeCall(cudaFree(sad));

	cutilSafeCall(cudaFreeArray(currentArray));
	cutilSafeCall(cudaFreeArray(lastArray));
}


void cuda_me(cuda_param_t *h, int* mvp[2], int* mv[2], int *mb_sad)
{
	cutilSafeCall(cudaMemcpy2DToArray(currentArray, 0, 0,
						h->fenc.plane, h->fenc.i_stride* sizeof(uint8_t),
						h->fenc.i_width * sizeof(uint8_t),  h->fenc.i_lines,
						cudaMemcpyHostToDevice));

	cutilSafeCall(cudaMemcpy2DToArray(lastArray, 0, 0,
						h->fref[0].plane,  h->fref[0].i_stride* sizeof(uint8_t),
						h->fref[0].i_width * sizeof(uint8_t), h->fref[0].i_lines,
						cudaMemcpyHostToDevice));

	cutilSafeCall(cudaBindTextureToArray(currentTex, 	currentArray, 	channelDesc));
	cutilSafeCall(cudaBindTextureToArray(lastTex, 		lastArray, 		channelDesc));

	int bw = (h->fenc.i_width + 15) / 16;
	int bh = (h->fenc.i_lines + 15) / 16;

    cudaMemcpy(mvXp, mvp[0], bw * bh * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(mvYp, mvp[1], bw * bh * sizeof(int), cudaMemcpyHostToDevice);

	runCuda(h->fenc.i_width, h->fenc.i_lines, mvXp, mvYp, mvX, mvY, sad);

    cudaMemcpy(mv[0], mvX, bh * bw * sizeof(int), cudaMemcpyDeviceToHost);
    cudaMemcpy(mv[1], mvY, bh * bw * sizeof(int), cudaMemcpyDeviceToHost);
    cudaMemcpy(mb_sad, sad, bh * bw * sizeof(int), cudaMemcpyDeviceToHost);

	cudaUnbindTexture(currentTex);
	cudaUnbindTexture(lastTex);
}

void runCuda(int width, int height, int *mvXp, int *mvYp, int *mvX, int *mvY, int *sad)
{
    int bw = (width + TILE - 1) / TILE;
    int bh = (height + TILE - 1) / TILE;

	dim3 blocks(bw, bh);
	dim3 threads(2 * SEARCHWINDOW + 1, 2 * SEARCHWINDOW + 1);

	minSAD<<<blocks, threads>>>(bw, mvXp, mvYp, mvX, mvY, sad);

}

__device__ unsigned mvCost(int dx, int dy){
	dx = (abs(dx))<<2;
	dy = (abs(dy))<<2;
	int xCost = round((log2f(dx+1)*2 + 0.718f + !!dx) + .5f);
	int yCost = round((log2f(dy+1)*2 + 0.718f + !!dy) + .5f);
	return (LAMBDA * (xCost+yCost));
}

__device__ unsigned SAD(ImageTexture a, ImageTexture b, unsigned xStart, unsigned yStart, unsigned height, unsigned width, unsigned dx, unsigned dy){
	unsigned ret = 0;
	for(unsigned x = xStart; x< xStart + width; x++){
		for(unsigned y = yStart; y< yStart + height; y++){
			ret = __usad(tex2D(a, x, y),tex2D(b, x + dx,y + dy), ret);
		}
	}
	return ret;
}

__global__ void minSAD(unsigned bw, int* mvXp, int* mvYp, int* mvX, int* mvY, int* sad)
{
	const unsigned int bx = blockIdx.x;
	const unsigned int by = blockIdx.y;

	const unsigned int tx = threadIdx.x;
	const unsigned int ty = threadIdx.y;

	const int blockTop = TILE * by;
	const int blockLeft = TILE * bx;

	const int dx = (tx - SEARCHWINDOW);
	const int dy = (ty - SEARCHWINDOW);

	__shared__  unsigned mins[(2 * SEARCHWINDOW + 1) * (2 * SEARCHWINDOW + 1)];
	__shared__  int2 smallestLoc[(2 * SEARCHWINDOW + 1) * (2 * SEARCHWINDOW + 1)];

	unsigned myLoc =tx + ty * (2 * SEARCHWINDOW + 1);

	int2 preMv;

    preMv.x = mvXp[bx + by * bw];
    preMv.y = mvYp[bx + by * bw];

	unsigned sum = SAD(currentTex, lastTex, blockLeft, blockTop, TILE, TILE, dx + preMv.x, dy + preMv.y);//

    sum += mvCost(dx, dy);
	mins[myLoc] = sum;

	smallestLoc[myLoc].x = dx + preMv.x;
	smallestLoc[myLoc].y = dy + preMv.y;

	__syncthreads();

	unsigned m=0;
	for(unsigned k = ((2 * SEARCHWINDOW + 1) * (2 * SEARCHWINDOW + 1)); k > 1; k = m){
		m = (k + 1) >> 1;
		if(myLoc < m && (m + myLoc < k))
		{
			if(mins[myLoc] > mins[myLoc + m]){
				smallestLoc[myLoc] = smallestLoc[myLoc + m];
				mins[myLoc] = mins[myLoc + m];
			}
		}

		__syncthreads();
	}

	if(myLoc == 0){
		mvX[bx + by * bw] = smallestLoc[0].x;
        mvY[bx + by * bw] = smallestLoc[0].y;
        sad[bx + by * bw] = mins[0];
	}
}

#ifdef __CUDACC__
}
#endif
